import librosa
import numpy as np
from sklearn.mixture import GaussianMixture
from scipy.spatial.distance import jensenshannon

import librosa
import numpy as np
from scipy import stats

def compute_kl(audio_path1, audio_path2):
    # 加载音频
    audio1, sr1 = librosa.load(audio_path1)
    audio2, sr2 = librosa.load(audio_path2)
    
    # 提取MFCC第一维
    mfcc1 = librosa.feature.mfcc(y=audio1, sr=sr1, n_mfcc=13)[0]  # 取第一维
    mfcc2 = librosa.feature.mfcc(y=audio2, sr=sr2, n_mfcc=13)[0]
    
    # 创建直方图分布
    bins = np.linspace(min(np.min(mfcc1), np.min(mfcc2)), max(np.max(mfcc1), np.max(mfcc2)), 100)
    hist1, _ = np.histogram(mfcc1, bins=bins, density=True)
    hist2, _ = np.histogram(mfcc2, bins=bins, density=True)
    
    # 归一化并平滑
    epsilon = 1e-10
    hist1 = hist1 + epsilon
    hist2 = hist2 + epsilon
    hist1 = hist1 / np.sum(hist1)
    hist2 = hist2 / np.sum(hist2)
    
    # 计算KL散度
    kl = np.sum(hist1 * np.log(hist1 / hist2))
    return kl

# 使用示例
ori = "Fireworks._High-quality,_stereo.wav"

js_value = compute_kl(ori, "Fireworks__High-qual_revert.wav")
print(f"Jensen–Shannon divergence: {js_value:.4f}")

js_value = compute_kl(ori, "Post_Rock__echoing_e_vae_test.wav")
print(f"Jensen–Shannon divergence: {js_value:.4f}")

js_value = compute_kl(ori, "Post_Rock__echoing_e_revert_1.15.wav")
print(f"Jensen–Shannon divergence: {js_value:.4f}")

# def compute_js_gmm(audio_path1, audio_path2, n_mfcc=13, n_components=8, n_samples=10000, random_state=42):
#     """
#     用多维 MFCC + 高斯混合模型 (GMM) + Jensen-Shannon 散度比较音频相似度
#     """
#     # 1. 加载音频
#     audio1, sr1 = librosa.load(audio_path1, sr=None)
#     audio2, sr2 = librosa.load(audio_path2, sr=None)

#     # 2. 提取 MFCC 特征（多维）
#     mfcc1 = librosa.feature.mfcc(y=audio1, sr=sr1, n_mfcc=n_mfcc).T  # shape: (frames, n_mfcc)
#     mfcc2 = librosa.feature.mfcc(y=audio2, sr=sr2, n_mfcc=n_mfcc).T

#     # 3. 拟合 GMM
#     gmm1 = GaussianMixture(n_components=n_components, covariance_type="diag", random_state=random_state).fit(mfcc1)
#     gmm2 = GaussianMixture(n_components=n_components, covariance_type="diag", random_state=random_state).fit(mfcc2)

#     # 4. 从两个 GMM 中采样
#     samples1, _ = gmm1.sample(n_samples)
#     samples2, _ = gmm2.sample(n_samples)

#     # 5. 估计概率分布（在联合样本上）
#     all_samples = np.vstack([samples1, samples2])
#     p1 = np.exp(gmm1.score_samples(all_samples))
#     p2 = np.exp(gmm2.score_samples(all_samples))

#     # 归一化为概率分布
#     p1 /= p1.sum()
#     p2 /= p2.sum()

#     # 6. 计算 Jensen-Shannon 散度
#     js = jensenshannon(p1, p2, base=2)  # 范围 [0, 1]
#     return js

# 使用示例
# ori = "Trance,_Ibiza,_Beach,_Sun,_4_AM,_Progressive,_Synthesizer,_909,_Dramatic_chords,_Choir,_Euphoric,_Nostalgic,_Dynamic,_Flowing.wav"

# js_value = compute_js_gmm(ori, "Trance__Ibiza__Beach_vae_test.wav")
# print(f"Jensen–Shannon divergence: {js_value:.4f}")

# js_value = compute_js_gmm(ori, "Trance__Ibiza__Beach_revert.wav")
# print(f"Jensen–Shannon divergence: {js_value:.4f}")

# js_value = compute_js_gmm(ori, "Trance__Ibiza__Beach_revert_1.15.wav")
# print(f"Jensen–Shannon divergence: {js_value:.4f}")